import numpy as np
from sklearn.model_selection import train_test_split
from sklearn.neighbors import KNeighborsClassifier
from sklearn.metrics import accuracy_score
import matplotlib.pyplot as plt
import joblib

# 准备数据
# 特征：年龄，每周空闲时间
X = np.array([[25, 10],
              [30, 12],
              [22, 5],
              [35, 8],
              [40, 15],
              [28, 7],
              [22, 10],
              [33, 14],
              [45, 12],
              [20, 20]])

# 标签：是否参加（1表示参加，0表示不参加）
y = np.array([1, 1, 0, 0, 1, 0, 1, 1, 0, 1])

# 数据分割
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)

# 模型选择
k = 3  # 选择K值
model = KNeighborsClassifier(n_neighbors=k)

# 模型训练
model.fit(X_train, y_train)

# 模型评估
y_pred = model.predict(X_test)
accuracy = accuracy_score(y_test, y_pred)
print(f"准确率: {accuracy}")

# 绘制决策边界
# 生成网格点
x_min, x_max = X[:, 0].min() - 1, X[:, 0].max() + 1
y_min, y_max = X[:, 1].min() - 1, X[:, 1].max() + 1
xx, yy = np.meshgrid(np.arange(x_min, x_max, 0.1),np.arange(y_min, y_max, 0.1))
# 对网格点进行预测
Z = model.predict(np.c_[xx.ravel(), yy.ravel()])
Z = Z.reshape(xx.shape)
# 绘制等高线
plt.contourf(xx, yy, Z, alpha=0.4)

# 绘制训练集样本点
plt.scatter(X_train[:, 0], X_train[:, 1], c=y_train, edgecolors='k', marker='o', label='Training data')

# 绘制测试集样本点
plt.scatter(X_test[:, 0], X_test[:, 1], c=y_test, marker='x', label='Testing data')
plt.xlabel('Age')
plt.ylabel('Hours of Free Time per Week')
plt.title('KNN Decision Boundary (Activity Participation)')
plt.legend()
plt.show()

# 保存模型
joblib_file = "knn_activity_model.pkl"
joblib.dump(model, joblib_file)
print(f"模型已保存到 {joblib_file}")

# 加载模型
loaded_model = joblib.load(joblib_file)

# 进行预测
new_data = np.array([[27, 11]])
prediction = loaded_model.predict(new_data)
print(f"预测值: {'参加' if prediction[0] == 1 else '不参加'}")
